# @Time : 2021/2/15 11:54
# @Author : Li Kunlun
# @Description :  进化策略
import numpy as np
import matplotlib.pyplot as plt

DNA_SIZE = 1  # DNA (real number)
DNA_BOUND = [0, 5]  # solution upper and lower bounds
N_GENERATIONS = 200
POP_SIZE = 100  # population size
N_KID = 50  # n kids per generation


def F(x): return np.sin(10 * x) * x + np.cos(2 * x) * x  # to find the maximum of this function


# find non-zero fitness for selection
def get_fitness(pred):
    return pred.flatten()


def make_kid(pop, n_kid):
    # generate empty kid holder（根据正态分布生孩子）
    kids = {'DNA': np.empty((n_kid, DNA_SIZE))}
    kids['mut_strength'] = np.empty_like(kids['DNA'])
    for kv, ks in zip(kids['DNA'], kids['mut_strength']):
        # crossover (roughly half p1 and half p2)
        # 随机选择父母进行交叉配对
        p1, p2 = np.random.choice(np.arange(POP_SIZE), size=2, replace=False)
        # p1=90([3.49778883]),p2=50([3.49778883])
        cp = np.random.randint(0, 2, DNA_SIZE, dtype=np.bool)  # crossover points
        # kv[cp] 和 kv[~cp]对应父母的DNA信息交叉
        kv[cp] = pop['DNA'][p1, cp]
        # 当cp对应的值为True时,将索引为p1的值赋值给kv
        #  kv[cp]--- [3.49778883]
        print(" kv[cp]---", kv[cp])
        kv[~cp] = pop['DNA'][p2, ~cp]
        #  kv[~cp]--- []
        print(" kv[~cp]---", kv[~cp])
        #  kv--- [3.49778883]
        print(" kv---", kv)
        # 将索引为p1的值赋值给ks,(kv,ks)对应的是同一个索引的字典值
        ks[cp] = pop['mut_strength'][p1, cp]
        ks[~cp] = pop['mut_strength'][p2, ~cp]

        # mutate (change DNA based on normal distribution)
        # 修改变异强度，选取一个最大值
        ks = np.maximum(ks + (np.random.rand(*ks.shape) - 0.5), 0.)  # must > 0
        # 在正态分布上取一个比较接近于均值的点
        kv += ks * np.random.randn(*kv.shape)
        # 修正坐标范围[0,5]之外的数据
        # kv需要添加上[:],不然会报错下面的警告
        # local variable kv value is not used
        kv[:] = np.clip(kv, *DNA_BOUND)  # clip the mutated value
    return kids


def kill_bad(pop, kids):
    # put pop and kids together（杀死坏孩子和坏父母）
    for key in ['DNA', 'mut_strength']:
        pop[key] = np.vstack((pop[key], kids[key]))

    fitness = get_fitness(F(pop['DNA']))  # calculate global fitness
    idx = np.arange(pop['DNA'].shape[0])
    # 从倒数第100个到最后都获取
    good_idx = idx[fitness.argsort()][-POP_SIZE:]  # selected by fitness ranking (not value)
    for key in ['DNA', 'mut_strength']:
        pop[key] = pop[key][good_idx]
    return pop


if __name__ == '__main__':
    pop = dict(DNA=5 * np.random.rand(1, DNA_SIZE).repeat(POP_SIZE, axis=0),  # initialize the pop DNA values
               mut_strength=np.random.rand(POP_SIZE, DNA_SIZE))  # initialize the pop mutation strength values

    plt.ion()
    for _ in range(N_GENERATIONS):
        plt.scatter(pop['DNA'], F(pop['DNA']), s=200, lw=0, c='red', alpha=0.5)
        x = np.linspace(*DNA_BOUND, 200)
        plt.plot(x, F(x))
        plt.pause(0.05)

        # ES part
        kids = make_kid(pop, N_KID)
        pop = kill_bad(pop, kids)  # keep some good parent for elitism

    plt.ioff()
    plt.show()
